import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions
from config import ChromaConfig
import pandas as pd
import logging
from tqdm import tqdm
import re
from typing import Dict, List, Any, Optional
from pathlib import Path

class ChromaManager:
    def __init__(self, config: ChromaConfig):
        self.config = config
        self.client = self._init_client()
        self.embedding_function = self._init_embedding_function()
        self.logger = logging.getLogger(__name__)

    def _init_client(self):
        if self.config.host:
            return chromadb.HttpClient(
                host=self.config.host,
                port=self.config.port
            )
        # print("Current Working Directory:", os.getcwd())
        # return chromadb.PersistentClient(path="./chroma_db")
        return chromadb.PersistentClient(path=str(Path("./chroma_db").resolve()))
        
    def _init_embedding_function(self):
        return embedding_functions.OllamaEmbeddingFunction(
            url=self.config.ollama_base_url + "/api/embeddings",
            model_name=self.config.embedding_model
        )
    
    def log(self, message: str, title: str = "Info"):
        print(f"{title}: {message}")
        
    def store_indicators(self, data: pd.DataFrame):
        total_rows = len(data)
        self.logger.info(f"Starting to store {total_rows} indicators...")
        
        collection = self.client.get_or_create_collection(
            name="indicators",
            embedding_function=self.embedding_function
        )
        
        success_count = 0
        error_count = 0
        
        for _, row in tqdm(data.iterrows(), total=total_rows, desc="Storing indicators"):
            try:
                collection.add(
                    documents=[row["指标名称"]],
                    metadatas=[{
                        "code": row["指标代码"],
                        "theme": row["主题描述"],
                        "formula": row["自定义计算表达式"]
                    }],
                    ids=[row["指标代码"]]
                )
                success_count += 1
            except Exception as e:
                self.logger.error(f"Error storing indicator {row['指标代码']}: {str(e)}")
                print(f"Error storing indicator {row['指标代码']}: {str(e)}")
                error_count += 1
                
        self.logger.info(f"""
        Storage Summary:
        - Total processed: {total_rows}
        - Successfully stored: {success_count}
        - Failed: {error_count}
        - Success rate: {(success_count/total_rows)*100:.2f}%
        """)
            
    def store_dimensions(self, data: pd.DataFrame):
        total_rows = len(data)
        self.logger.info(f"Starting to store {total_rows} dimensions...")
        
        collection = self.client.get_or_create_collection(
            name="dimensions",
            embedding_function=self.embedding_function
        )
        
        success_count = 0
        error_count = 0
        
        for _, row in tqdm(data.iterrows(), total=total_rows, desc="Storing dimensions"):
            try:
                collection.add(
                    documents=[row["维度描述"]],
                    metadatas=[{
                        "code": row["维度代码"],
                        "theme": row["主题描述"]
                    }],
                    ids=[f"{row['维度代码']}_{row['主题描述']}"]
                )
                success_count += 1
            except Exception as e:
                self.logger.error(f"Error storing dimension {row['维度代码']}: {str(e)}")
                print(f"Error storing dimension {row['维度代码']}: {str(e)}")
                error_count += 1
                
        self.logger.info(f"""
        Storage Summary:
        - Total processed: {total_rows}
        - Successfully stored: {success_count}
        - Failed: {error_count}
        - Success rate: {(success_count/total_rows)*100:.2f}%
        """)

    def query_indicators(self, query_text: str, limit: int = 10, include_related: bool = False) -> dict:
        """Query indicators collection by text similarity
        
        Args:
            query_text: Text to search for
            limit: Maximum number of results
            include_related: Include related indicators from formula
            
        Returns:
            Dict with 'indicators' and 'themes' lists
        """
        try: 
            collection = self.client.get_collection(
                name="indicators",
                embedding_function=self.embedding_function
            )

            print(f"Querying for '{query_text}'...")
            
            try: 
                results = collection.query(
                    query_texts=[query_text],
                    n_results=limit
                )
            except Exception as e:
                print("Error occurred during query:", e)

            print(f"Found matches")
            
            all_indicators = []
            themes = set()  # Use set for unique themes

            for i in range(len(results['ids'][0])):
                match = {
                    'id': results['ids'][0][i],
                    'text': results['documents'][0][i],
                    'metadata': results['metadatas'][0][i],
                    'distance': results['distances'][0][i]
                }
                all_indicators.append(match)
                theme = match['metadata'].get('theme')
                if theme:
                    themes.add(theme)

                # Query related indicators if requested
                if include_related and 'formula' in match['metadata']:
                    formula = match['metadata']['formula']
                    related_codes = self._parse_formula(formula)

                    for code in related_codes:
                        related_result = collection.get([code])
                        # Check if related_result and its ids exist and are not empty
                        if (related_result 
                            and 'ids' in related_result 
                            and related_result['ids'] 
                            and len(related_result['ids']) > 0):
                            related = {
                                'id': related_result['ids'][0],
                                'text': related_result['documents'][0],
                                'metadata': related_result['metadatas'][0],
                                'distance': 0,  # Direct lookup, no distance
                            }
                            all_indicators.append(related)
                            theme = related['metadata'].get('theme')
                            if theme:
                                themes.add(theme)

            # Filter by threshold
            # return [m for m in matches if m['distance'] <= threshold]
            # Deduplicate and sort by distance
            deduped = self._deduplicate_indicators(all_indicators)

            return {
                'indicators': sorted(deduped, key=lambda x: x['distance']),
                'themes': sorted(list(themes))
            }
        except Exception as e:
            self.log(str(e), "Error generating KPI params")
            raise
    
    def query_dimensions(self, query_text: str, limit: int = 10, themes: List[str] = None) -> list:
        """Query dimensions collection by text similarity
        
        Args:
            query_text: Text to search for
            limit: Maximum number of results
            themes: Optional list of themes to filter by
            
        Returns:
            List of matching dimensions with metadata
        """
        collection = self.client.get_collection(
            name="dimensions",
            embedding_function=self.embedding_function
        )

        # Build where clause if themes provided
        where = {"theme": {"$in": themes}} if themes else None
        
        results = collection.query(
            query_texts=[query_text],
            n_results=limit,
            where=where
        )
        
        matches = []
        for i in range(len(results['ids'][0])):
            matches.append({
                'id': results['ids'][0][i], 
                'text': results['documents'][0][i],
                'metadata': results['metadatas'][0][i],
                'distance': results['distances'][0][i]
            })
            
        # Filter by threshold
        # return [m for m in matches if m['distance'] <= threshold]
        return matches

    def _generate_markdown_table(self, headers: list, rows: list) -> str:
        """Generate markdown table from headers and rows"""
        # Headers row
        output = "| " + " | ".join(headers) + " |\n"
        # Separator row
        output += "|" + "|".join(["---" for _ in headers]) + "|\n"
        # Data rows
        for row in rows:
            output += "| " + " | ".join(str(cell) for cell in row) + " |\n"
        return output

    def get_indicators_table(self, query_text: str, limit: int = 10, include_related: bool = False) -> str:
        """Query indicators and return markdown table"""
        try:
            results = self.query_indicators(query_text, limit, include_related)
            return self.format_indicators_table(results['indicators'])
        except Exception as e:
            self.logger.error(f"Error generating indicators table: {str(e)}")
            return f"Error: {str(e)}"
        
    def format_indicators_table(self, indicators: list) -> str:
        """Query indicators and return markdown table"""
        try:
            headers = ["指标代码", "指标描述", "关联主题", "计算公式"]
            rows = []
            for item in indicators:
                rows.append([
                    item['metadata']['code'],
                    item['text'],
                    item['metadata'].get('theme','-'),
                    item['metadata'].get('formula', '-')
                ])
            return self._generate_markdown_table(headers, rows)
        except Exception as e:
            self.logger.error(f"Error formatting indicators table: {str(e)}")
            return f"Error: {str(e)}"

    def get_dimensions_table(self, query_text: str, limit: int = 10, themes: List[str] = None) -> str:
        """Query dimensions and return markdown table
        
        Args:
            query_text: Text to search for
            limit: Maximum number of results
            themes: Optional list of themes to filter by
        """
        try:
            results = self.query_dimensions(query_text, limit, themes=themes)

            # # Filter by themes if provided
            # if themes:
            #     results = [
            #         item for item in results 
            #         if item['metadata']['theme'] in themes
            #     ]

            headers = ["维度代码", "维度描述", "所属主题"]
            rows = []
            for item in results:
                rows.append([
                    item['metadata']['code'],
                    item['text'],
                    item['metadata']['theme']
                ])
            return self._generate_markdown_table(headers, rows)
        except Exception as e:
            self.logger.error(f"Error generating dimensions table: {str(e)}")
            return f"Error: {str(e)}"
        
    def _parse_formula(self, formula: str) -> list:
        """Extract indicator codes from formula
        
        Args:
            formula: Formula string like "{S0010078/S0010018}" or "{S0010078+S0010019}"
            
        Returns:
            List of unique indicator codes
        """
        if not formula:
            return []
            
        # Remove curly braces
        formula = formula.strip('{}')
        
        # Use regex to split by operators
        codes = re.split(r'[+\-*/]', formula)
        
        # Remove duplicates and empty strings
        return list(set(filter(None, codes)))
    
    def _deduplicate_indicators(self, indicators: list) -> list:
        """Deduplicate indicators based on metadata.code"""
        seen = {}
        for item in indicators:
            code = item['metadata']['code']
            if code not in seen:
                seen[code] = item
        return list(seen.values())